#!/usr/bin/env python
# Copyright (C) 2007  "Peter Roesch" <Peter.Roesch@fh-augsburg.de>
#
# This code is licensed under the PyOpenGL License.
# Details are given in the file license.txt included in this distribution.
# import OpenGL
#OpenGL.FULL_LOGGING = True

import sys
import array
import random

import PIL.Image as Image

try:
    from OpenGL.GLUT import *
    from OpenGL.GL import *
    from OpenGL.GLU import *
    from OpenGL.GL.ARB.shader_objects import *
    from OpenGL.GL.ARB.fragment_shader import *
    from OpenGL.GL.ARB.vertex_shader import *
except ImportError:
    print(''' Error PyOpenGL not installed properly !!''')
    sys.exit()


class ShaderProgram(object):
    """Manage GLSL programs."""

    def __init__(self):
        self.__requiredExtensions = ["GL_ARB_fragment_shader", "GL_ARB_vertex_shader", "GL_ARB_shader_objects",
                                     "GL_ARB_shading_language_100", "GL_ARB_vertex_shader", "GL_ARB_fragment_shader"]
        self.checkExtensions(self.__requiredExtensions)
        self.__shaderProgramID = glCreateProgramObjectARB()
        self.__checkOpenGLError()
        self.__programReady = False
        self.__isEnabled = False
        self.__shaderObjectList = []

    def checkExtensions(self, extensions):
        """Check if all extensions in a list are present."""
        for ext in extensions:
            if ( not ext ):
                print("Driver does not support ", ext)
                sys.exit()

    def __checkOpenGLError(self):
        """Print OpenGL error message."""
        err = glGetError()
        if ( err != GL_NO_ERROR ):
            print('GLERROR: ', gluErrorString(err))
            sys.exit()

    def reset(self):
        """Disable and remove all shader programs"""
        for shaderID in self.__shaderObjectList:
            glDetachObjectARB(self.__shaderProgramID, shaderID)
            glDeleteObjectARB(shaderID)
            self.__shaderObjectList.remove(shaderID)
            self.__checkOpenGLError()
        glDeleteObjectARB(self.__shaderProgramID)
        self.__checkOpenGLError()
        self.__shaderProgramID = glCreateProgramObjectARB()
        self.__checkOpenGLError()
        self.__programReady = False

    def addShader(self, shaderType, fileName):
        """Read a shader program from a file.

        The program is load and compiled"""
        shaderHandle = glCreateShaderObjectARB(shaderType)
        self.__checkOpenGLError()
        sourceString = open(fileName, 'r').read()
        glShaderSourceARB(shaderHandle, [sourceString])
        self.__checkOpenGLError()
        glCompileShaderARB(shaderHandle)
        success = glGetObjectParameterivARB(shaderHandle, GL_OBJECT_COMPILE_STATUS_ARB)
        if (not success):
            print(glGetInfoLogARB(shaderHandle))
            sys.exit()
        glAttachObjectARB(self.__shaderProgramID, shaderHandle)
        self.__checkOpenGLError()
        self.__shaderObjectList.append(shaderHandle)

    def linkShaders(self):
        """Link compiled shader programs."""
        glLinkProgramARB(self.__shaderProgramID)
        self.__checkOpenGLError()
        success = glGetObjectParameterivARB(self.__shaderProgramID, GL_OBJECT_LINK_STATUS_ARB)
        if (not success):
            print(glGetInfoLogARB(self.__shaderProgramID))
            sys.exit()
        else:
            self.__programReady = True

    def enable(self):
        """Activate shader programs."""
        if self.__programReady:
            glUseProgramObjectARB(self.__shaderProgramID)
            self.__isEnabled = True
            self.__checkOpenGLError()
        else:
            print("Shaders not compiled/linked properly, enable() failed")

    def disable(self):
        """De-activate shader programs."""
        glUseProgramObjectARB(0)
        self.__isEnabled = False
        self.__checkOpenGLError()

    def indexOfUniformVariable(self, variableName):
        """Find the index of a uniform variable."""
        if not self.__programReady:
            print("\nShaders not compiled/linked properly")
            result = -1
        else:
            result = glGetUniformLocationARB(self.__shaderProgramID, variableName)
            self.__checkOpenGLError()
        if result < 0:
            print('Variable "%s" not known to the shader' % ( variableName ))
            sys.exit()
        else:
            return result

    def indexOfVertexAttribute(self, attributeName):
        """Find the index of an attribute variable."""
        if not self.__programReady:
            print("\nShaders not compiled/linked properly")
            result = -1
        else:
            result = glGetAttribLocationARB(self.__shaderProgramID, attributeName)
            self.__checkOpenGLError()
        if result < 0:
            print('Attribute "%s" not known to the shader' % ( attributeName ))
            sys.exit()
        else:
            return result

    def isEnabled(self):
        return self.__isEnabled


class Texture(object):
    def __init__(self):
        self.xSize, self.ySize = 0, 0
        self.rawRefence = None


class RandomTexture(Texture):
    def __init__(self, xSizeP, ySizeP):
        Texture.__init__(self)
        self.xSize, self.ySize = xSizeP, ySizeP
        tmpList = [random.randint(0, 255) for i in range(3 * self.xSize * self.ySize)]
        self.textureArray = array.array('B', tmpList)
        self.rawReference = self.textureArray.tostring()


class FileTexture(Texture):
    def __init__(self, fileName):
        Texture.__init__(self)
        im = Image.open(fileName)
        self.xSize = im.size[0]
        self.ySize = im.size[1]
        self.rawReference = im.tostring("raw", "RGB", 0, -1)


frameRate = 30
from time import sleep


def animationStep(*args):
    global frameRate
    global sP
    if not quadList:
        if len(sys.argv) > 1:
            init(sys.argv[1])
        else:
            init(None)
        assert quadList
    if sP.isEnabled():
        global rgbTransformMatrix
        row = random.randint(0, 3)
        column = random.randint(0, 3)
        rgbTransformMatrix[row * 4 + column] += random.normalvariate(0, 0.05)
        if rgbTransformMatrix[row * 4 + column] < 0:
            rgbTransformMatrix[row * 4 + column] = 0
        sum = 0.
        for x in rgbTransformMatrix[row * 4:row * 4 + 4]:
            sum += x
        for i in range(row * 4, row * 4 + 4):
            rgbTransformMatrix[i] /= sum
        glUniformMatrix4fvARB(sP.indexOfUniformVariable("RGBTransformationMatrix"), 1, False, rgbTransformMatrix)
    sleep(1 / float(frameRate))
    glutPostRedisplay()


def display(*args):
    if not quadList:
        if len(sys.argv) > 1:
            init(sys.argv[1])
        else:
            init(None)
        assert quadList
    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
    glColor3f(1, 1, 1)
    glMatrixMode(GL_PROJECTION)
    glLoadIdentity()
    xSize, ySize = glutGet(GLUT_WINDOW_WIDTH), glutGet(GLUT_WINDOW_HEIGHT)
    f = float(ySize) / float(xSize)
    gluOrtho2D(-1.5, 1.5, -1.5 * f, 1.5 * f)
    glMatrixMode(GL_MODELVIEW)
    glLoadIdentity()
    glPushMatrix()
    glTranslatef(-0.6, 0, 0)
    sP.disable()
    glCallList(quadList)
    glPopMatrix()
    glTranslatef(0.6, 0, 0)
    sP.enable()
    glCallList(quadList)
    glutSwapBuffers()


sP = None


def initShaders():
    global sP
    sP = ShaderProgram()
    sP.addShader(GL_FRAGMENT_SHADER_ARB, "rgbMorph.frag")
    sP.linkShaders()
    sP.enable()
    global rgbTransformMatrix
    glUniformMatrix4fvARB(sP.indexOfUniformVariable("RGBTransformationMatrix"), 1, False, rgbTransformMatrix)


quadList = None
rgbTransformMatrix = [0.0 for i in range(16)]
rgbTransformMatrix[0] = rgbTransformMatrix[5] = rgbTransformMatrix[10] = rgbTransformMatrix[15] = 1.0


def init(fileName):
    try:
        texture = FileTexture(fileName)
    except:
        print('could not open ', fileName, '; using random texture')
        texture = RandomTexture(256, 256)
    glClearColor(0, 0, 0, 0)
    glShadeModel(GL_SMOOTH)
    glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_REPEAT)
    glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_REPEAT)
    glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
    glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR)
    glTexImage2D(GL_TEXTURE_2D, 0, 3, texture.xSize, texture.ySize, 0, GL_RGB, GL_UNSIGNED_BYTE, texture.rawReference)
    glEnable(GL_TEXTURE_2D)
    global quadList
    quadList = glGenLists(1)
    glNewList(quadList, GL_COMPILE)
    glBegin(GL_QUADS)
    glTexCoord2f(0, 1)
    glVertex3f(-0.5, 0.5, 0)
    glTexCoord2f(0, 0)
    glVertex3f(-0.5, -0.5, 0)
    glTexCoord2f(1, 0)
    glVertex3f(0.5, -0.5, 0)
    glTexCoord2f(1, 1)
    glVertex3f(0.5, 0.5, 0)
    glEnd()
    glEndList()
    initShaders()


def main():
    import logging

    logging.basicConfig()
    glutInit(sys.argv)
    glutInitDisplayMode(GLUT_DOUBLE | GLUT_RGB)
    glutInitWindowSize(512, 512)
    glutInitWindowPosition(100, 100)
    glutCreateWindow(sys.argv[0])
    glutDisplayFunc(display)
    glutIdleFunc(animationStep)
    glutMainLoop()


if __name__ == "__main__":
    main()

